package com.gitee.wsl.ext.collection

import com.gitee.wsl.ext.K
import com.gitee.wsl.ext.runBlocking
import kotlinx.coroutines.CoroutineStart
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.async
import kotlinx.coroutines.withContext
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.EmptyCoroutineContext


/**
 * Execute [block] in parallel using operator async for each element of the collection
 */
suspend fun <T> Iterable<T>.forEachParallel(
    context: CoroutineContext = EmptyCoroutineContext,
    start: CoroutineStart = CoroutineStart.DEFAULT,
    block: suspend (T) -> Unit,
): Unit = withContext(context) {
    map { async(context, start) { block(it) } }.forEach { it.await() }
}

/**
 * Execute [block] in parallel using operator async for each element of the collection
 */
suspend fun <T, R> Iterable<T>.mapParallel(
    context: CoroutineContext = EmptyCoroutineContext,
    start: CoroutineStart = CoroutineStart.DEFAULT,
    block: suspend (T) -> R,
): List<R> = withContext(context) {
    map { async(context, start) { block(it) } }.map { it.await() }
}

//
// Parallel collection utilities
//

fun <A, B> Iterable<A>.asyncMap(
//    numThreads: Int = Runtime.getRuntime().availableProcessors(),
    f: suspend (A) -> B
): List<B> = K.runBlocking {
    map { async(Dispatchers.Default) { f(it) } }.map { it.await() }
}

@Suppress("unused")
fun <A, B> Iterable<A>.asyncEach(
//    numThreads: Int = Runtime.getRuntime().availableProcessors(),
    f: suspend (A) -> B
) {
    asyncMap(f).count() // count is just needed to force computations
}
